package me.luraframework.gateway.module.log;

import com.google.common.base.Strings;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import me.luraframework.gateway.utils.IpAddrUtils;
import net.dreamlu.mica.ip2region.core.Ip2regionSearcher;
import net.dreamlu.mica.ip2region.core.IpInfo;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.StopWatch;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.zip.GZIPInputStream;

import static org.springframework.cloud.gateway.filter.NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER;

/**
 * @ description: AccessLogFilter
 * @ author: Liu Ran
 * @ data: 7/11/23 10:47
 */

@Slf4j
@Component
@RequiredArgsConstructor
public class AccessLogFilter implements GlobalFilter, Ordered {

    private final GatewayLogRepository gatewayLogRepository;
    private final Ip2regionSearcher ip2regionSearcher;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();

        // 请求路径
        String requestPath = request.getPath().pathWithinApplication().value();

        Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);


        String ipAddress = IpAddrUtils.getIp(request);
        StopWatch stopWatch = new StopWatch();

        GatewayLog gatewayLog = new GatewayLog();
        gatewayLog.setUrl(request.getURI().toString());
        gatewayLog.setRequestMethod(request.getMethodValue());
        gatewayLog.setRequestPath(requestPath);
        gatewayLog.setTargetServer(route.getId());
        gatewayLog.setRequestTime(LocalDateTime.now());
        gatewayLog.setIp(ipAddress);
        stopWatch.start();
        CachedBodyOutputMessage cachedBodyOutputMessage = new CachedBodyOutputMessage(exchange, exchange.getRequest().getHeaders());


        return cachedBodyOutputMessage.writeWith(exchange.getRequest().getBody())
                .and(chain.filter(exchange.mutate().request(new ServerHttpRequestDecorator(exchange.getRequest()) {

                            @Override
                            public Flux<DataBuffer> getBody() {
                                return cachedBodyOutputMessage.getBody();
                            }
                        }).response(serverHttpResponseDecorator(exchange, gatewayLog)).build())
                        .doFinally(f -> writeAccessLog(gatewayLog, stopWatch)));


    }

    private ServerHttpResponseDecorator serverHttpResponseDecorator(ServerWebExchange exchange, GatewayLog gatewayLog) {
        return new ServerHttpResponseDecorator(exchange.getResponse()) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                    // DataBufferFactory合并多个dataBuffers
                    DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                    DataBuffer join = dataBufferFactory.join(dataBuffers);
                    byte[] content = new byte[join.readableByteCount()];
                    join.read(content);

                    DataBufferUtils.release(join);

                    String acceptEncoding = exchange.getRequest().getHeaders().getFirst(HttpHeaders.ACCEPT_ENCODING);
                    String responseResult;
                    if (!Strings.isNullOrEmpty(acceptEncoding) && acceptEncoding.contains("gzip")) {
                        try {
                            responseResult = new String(uncompress(content), StandardCharsets.UTF_8);
                        } catch (Exception e) {
                            responseResult = new String(content, StandardCharsets.UTF_8);
                        }
                    } else {
                        responseResult = new String(content, StandardCharsets.UTF_8);
                    }

                    gatewayLog.setResponseData(responseResult);
                    gatewayLog.setResponseSize(content.length);
                    return exchange.getResponse().bufferFactory().wrap(content);
                }));
            }
        };
    }

    private void writeAccessLog(GatewayLog gatewayLog, StopWatch stopWatch) {
        stopWatch.stop();
        gatewayLog.setAddress(getCityInfo(gatewayLog.getIp()));
        gatewayLog.setResponseTime(LocalDateTime.now());
        gatewayLog.setExecuteTime(stopWatch.getTotalTimeMillis());
        gatewayLogRepository.save(gatewayLog).subscribe();
    }

    @Override
    public int getOrder() {
        return WRITE_RESPONSE_FILTER_ORDER ;
    }

    public static byte[] uncompress(byte[] bytes) {
        if (bytes == null || bytes.length == 0) {
            return null;
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ByteArrayInputStream in = new ByteArrayInputStream(bytes);
        try {
            GZIPInputStream ungzip = new GZIPInputStream(in);
            byte[] buffer = new byte[256];
            int n;
            while ((n = ungzip.read(buffer)) >= 0) {
                out.write(buffer, 0, n);
            }
        } catch (IOException e) {
            log.error("gzip uncompress error.", e);
        }

        return out.toByteArray();

    }

    public String getCityInfo(String ip) {
        IpInfo ipInfo = ip2regionSearcher.memorySearch(ip);
        if(ipInfo != null){
            return ipInfo.getAddress();
        }
        return null;
    }
}
